from thop import profile
import os
from thop import clever_format

from ana import buildMambaIR
import torch
import torch.nn as nn
import torch.nn.functional as F
import time

if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    H=128
    W=128
    scale=4
    # NOTE: the default returned model is classic SRx4 model
    init_model = buildMambaIR(upscale=scale).to(device)
    dummy_input = torch.randn(1, 3, H//scale,W//scale).to(device)
    flops, params = profile(init_model, inputs=(dummy_input,))
    flops, params = clever_format([flops, params], '%.3f')
    print(flops)
    print(params)
    import torch
    from mamba_ssm import Mamba

    batch, length, dim = 2, 64, 16
    x = torch.randn(batch, length, dim).to("cuda")
    model = Mamba(
        # This module uses roughly 3 * expand * d_model^2 parameters
        d_model=dim, # Model dimension d_model
        d_state=16,  # SSM state expansion factor
        d_conv=4,    # Local convolution width
        expand=2,    # Block expansion factor
    ).to("cuda")
    y = model(x)
    assert y.shape == x.shape
    flops, params = profile(model, inputs=(x,))
    flops, params = clever_format([flops, params], '%.3f')
    print(flops)
    print(params)
    from calflops import calculate_flops
    flops, macs, params = calculate_flops(model=init_model, 
                                        input_shape=tuple(dummy_input.shape),
                                        output_as_string=True,
                                        output_precision=4)
    print("Alexnet FLOPs:%s   MACs:%s   Params:%s \n" %(flops, macs, params))